{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJcYs_ERTnnI"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HMUDt0CiUJk9"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "77z2OchJTk0l"
      },
      "source": [
        "# Migrating your TPUEstimator to TF2\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/migrate/tpu_estimator\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/tpu_estimator.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/tpu_estimator.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/tpu_estimator.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "meUTrR4I6m1C"
      },
      "source": [
        "In TF1, `tf.compat.v1.estimator.tpu.TPUEstimator` is a high level API that encapsulates training, evaluation, prediction, and export for serving with TPUs. In TF2, use `tf.keras` with a Keras [model](https://www.tensorflow.org/api_docs/python/tf/keras/Model?version=nightly), [layers](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer?version=nightly), [optimizers](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Optimizer?version=nightly), and [metrics](https://www.tensorflow.org/api_docs/python/tf/keras/metrics/Metric) for the aforementioned tasks. In this guide, you will go through the basic setup of a `TPUEstimator` training and evaluation program to run it. Then, you will create the equivalent training and evaluation program in TF2 using Keras, with the help of `tf.distribute.TPUStrategy`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdZSoIXEbhg-"
      },
      "source": [
        "## Setup\n",
        "\n",
        "For both TF1 and TF2 examples demonstrated below, you will first start with a couple of necessary TensorFlow imports,"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iE0vSfMXumKI"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jsm9Rxx7s1OZ"
      },
      "source": [
        "and prepare some simple data for demonstration:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m7rnGxsXtDkV"
      },
      "outputs": [],
      "source": [
        "features = [[1., 1.5]]\n",
        "labels = [[0.3]]\n",
        "eval_features = [[4., 4.5]]\n",
        "eval_labels = [[0.8]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4uXff1BEssdE"
      },
      "source": [
        "## TF1: TPUEstimator.train/evaluate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BVWHEQj5a7rN"
      },
      "source": [
        "Now, to use a `TPUEstimator`, you should first define a few functions: an input function for the training data, an evaluation input function for the evaluation data, and a model function that tells the `TPUEstimator` how the training op is defined with the features and labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lqe9obf7suIj"
      },
      "outputs": [],
      "source": [
        "def _input_fn(params):\n",
        "  dataset = tf1.data.Dataset.from_tensor_slices((features, labels))\n",
        "  dataset = dataset.repeat()\n",
        "  return dataset.batch(params['batch_size'], drop_remainder=True)\n",
        "\n",
        "def _eval_input_fn(params):\n",
        "  dataset = tf1.data.Dataset.from_tensor_slices((eval_features, eval_labels))\n",
        "  dataset = dataset.repeat()\n",
        "  return dataset.batch(params['batch_size'], drop_remainder=True)\n",
        "\n",
        "def _model_fn(features, labels, mode, params):\n",
        "  logits = tf1.layers.Dense(1)(features)\n",
        "  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)\n",
        "  optimizer = tf1.train.AdagradOptimizer(0.05)\n",
        "  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())\n",
        "  return tf1.estimator.tpu.TPUEstimatorSpec(mode, loss=loss, train_op=train_op)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QYnP3Dszc-2R"
      },
      "source": [
        "With those functions defined, you will now create a `tf.distribute.cluster_resolver.TPUClusterResolver` that provides the cluster information, and a `tf.compat.v1.estimator.tpu.RunConfig` object. Along with the model function you have defined, you can now create a `TPUEstimator`. Note that here you will simplify the flow by skipping checkpoint savings, and need to specify the batch size for both training and evaluation for a `TPUEstimator`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WAqyqawemlcl"
      },
      "outputs": [],
      "source": [
        "cluster_resolver = tf1.distribute.cluster_resolver.TPUClusterResolver(tpu='')\n",
        "print(\"All devices: \", tf1.config.list_logical_devices('TPU'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HsOpjW5plH9Q"
      },
      "outputs": [],
      "source": [
        "tpu_config = tf1.estimator.tpu.TPUConfig(iterations_per_loop=10)\n",
        "config = tf1.estimator.tpu.RunConfig(\n",
        "    cluster=cluster_resolver,\n",
        "    save_checkpoints_steps=None,\n",
        "    tpu_config=tpu_config)\n",
        "estimator = tf1.estimator.tpu.TPUEstimator(\n",
        "    model_fn=_model_fn, config=config, train_batch_size=8, eval_batch_size=8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uxw7tWrcepaZ"
      },
      "source": [
        "Call `train` to train the `TPUEstimator`,"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WZPKFOMAcyrP"
      },
      "outputs": [],
      "source": [
        "estimator.train(_input_fn, steps=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ev1vjIz9euIw"
      },
      "source": [
        "And call `evaluate` to evaluate the model using the evaluation data you have prepared."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bqiKRiwWc0cz"
      },
      "outputs": [],
      "source": [
        "estimator.evaluate(_eval_input_fn, steps=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEmzBjfnsxwT"
      },
      "source": [
        "### TF2: Keras training API with `tf.distribute.TPUStrategy`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UesuXNbShrbi"
      },
      "source": [
        "In TF2, you will use a `tf.keras` model with a `tf.distribute.TPUStrategy` to train your models on the TPU workers. To begin with, you will create a `TPUClusterResolver` to provide the cluster information, and connect to the cluster."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_TgdPNgXoS63"
      },
      "outputs": [],
      "source": [
        "cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')\n",
        "tf.config.experimental_connect_to_cluster(cluster_resolver)\n",
        "tf.tpu.experimental.initialize_tpu_system(cluster_resolver)\n",
        "print(\"All devices: \", tf.config.list_logical_devices('TPU'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R4EHXhN3CVmo"
      },
      "source": [
        "Next, once your data is prepared, you will create a `TPUStrategy`, and define your model, metrics, and optimizer under the scope of this strategy. You should pick a number for `steps_per_execution` in `Model.compile` because it specifies the number of batches to run during each `tf.function` call, and is critical for performance. This argument is similar to `iterations_per_loop` used in `TPUEstimator`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "atVciNgPs0fw"
      },
      "outputs": [],
      "source": [
        "dataset = tf.data.Dataset.from_tensor_slices(\n",
        "    (features, labels)).shuffle(10).repeat().batch(8).prefetch(2)\n",
        "eval_dataset = tf.data.Dataset.from_tensor_slices(\n",
        "      (eval_features, eval_labels)).batch(1)\n",
        "\n",
        "strategy = tf.distribute.TPUStrategy(cluster_resolver)\n",
        "with strategy.scope():\n",
        "  model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])\n",
        "  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)\n",
        "  model.compile(optimizer, \"mse\", steps_per_execution=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FkM2VZyni98F"
      },
      "source": [
        "With that, you are ready to train the model with the training dataset,"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kip65sYBlKiu"
      },
      "outputs": [],
      "source": [
        "model.fit(dataset, epochs=5, steps_per_epoch=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r0AEK8sNjLOj"
      },
      "source": [
        "and evaluate the model using the evaluation dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6tMRkyfKhqSL"
      },
      "outputs": [],
      "source": [
        "model.evaluate(eval_dataset, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gHx_RUL8xcJ3"
      },
      "source": [
        "For more customization such as optimization procedure, Keras APIs allow you to define the training step. Please see [Customize what happens in Model.fit](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit\n",
        ") for an example of this workflow, and [Use TPUs](https://www.tensorflow.org/guide/tpu) for more information of TPU usage, including an example of customized training loop."
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "name": "tpu_estimator.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
